// visualizer/src/App.js
import React, { useState, useEffect, useRef } from 'react';
import './App.css';

function App() {
  const [gameStatus, setGameStatus] = useState({
    cycle: 0,
    iron_curtain: 0,
    training_status: "已停止",
    factors_count: 12,
    active_connections: 0
  });
  const [trainingStatus, setTrainingStatus] = useState({
    is_running: false,
    current_cycle: 0,
    total_cycles: 0,
    progress: 0.0
  });
  const [activeFactors, setActiveFactors] = useState([]);
  const [trainingMetrics, setTrainingMetrics] = useState({
    overall_performance: 0.0,
    average_reward: 0.0,
    survival_rate: 0.0,
    training_loss: 0.0
  });
  const [logs, setLogs] = useState([]);
  const [isTraining, setIsTraining] = useState(false);
  const ws = useRef(null);

  // WebSocket连接 - 主要数据源
  useEffect(() => {
    const connectWebSocket = () => {
      const wsUrl = `ws://localhost:8001/ws`;
      
      console.log('正在连接到WebSocket:', wsUrl);
      
      ws.current = new WebSocket(wsUrl);
      
      ws.current.onopen = () => {
        console.log('WebSocket连接已建立');
        addLog('✅ WebSocket连接已建立 - 实时数据开始推送');
      };
      
      ws.current.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data);
          console.log('收到WebSocket消息:', data);
          handleWebSocketMessage(data);
        } catch (error) {
          console.error('解析WebSocket消息失败:', error);
        }
      };
      
      ws.current.onclose = () => {
        console.log('WebSocket连接断开，3秒后重连...');
        addLog('⚠️ WebSocket连接断开，3秒后重连...');
        setTimeout(connectWebSocket, 3000);
      };
      
      ws.current.onerror = (error) => {
        console.error('WebSocket错误:', error);
        addLog(`❌ WebSocket连接错误: ${error.message}`);
      };
    };

    connectWebSocket();

    return () => {
      if (ws.current) {
        ws.current.close();
      }
    };
  }, []);

  const handleWebSocketMessage = (data) => {
    switch(data.type) {
      case 'system_status':
        // 直接从WebSocket更新游戏状态
        setGameStatus(prev => ({
          ...prev,
          cycle: data.cycle,
          iron_curtain: data.iron_curtain,
          training_status: data.training_status,
          active_connections: data.active_connections
        }));
        break;
        
      case 'factor_update':
        setActiveFactors(data.factors);
        break;
        
      case 'training_metrics':
        setTrainingMetrics(data);
        break;
        
      default:
        console.log('未知消息类型:', data.type);
    }
  };

  // 只获取训练状态（不获取其他数据，避免冲突）
  const fetchTrainingStatus = async () => {
    try {
      const response = await fetch('http://localhost:8000/api/training/status');
      const data = await response.json();
      
      setTrainingStatus({
        is_running: data.is_running,
        progress: data.progress,
        current_cycle: data.current_cycle,
        total_cycles: data.total_cycles
      });
      
      setIsTraining(data.is_running);
      
    } catch (error) {
      console.error('获取训练状态失败:', error);
    }
  };

  // 简化定期获取，只获取训练状态
  useEffect(() => {
    const interval = setInterval(() => {
      fetchTrainingStatus();
    }, 5000); // 5秒获取一次训练状态
    
    // 立即执行一次
    fetchTrainingStatus();
    
    return () => clearInterval(interval);
  }, []);

  const addLog = (message) => {
    setLogs(prev => {
      // 避免重复日志
      if (prev.length > 0 && prev[prev.length - 1].message === message) {
        return prev;
      }
      return [...prev.slice(-19), {
        time: new Date().toLocaleTimeString(),
        message
      }];
    });
  };

  // 控制函数
  const startTraining = async () => {
    try {
      addLog('🚀 正在启动训练...');
      
      const response = await fetch('http://localhost:8000/api/training/start', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({
          cycles: 10,
          episodes: 5,
          model_name: "lightweight_ppo",
          learning_rate: 0.0003,
          batch_size: 64
        })
      });
      
      const result = await response.json();
      addLog(`✅ 训练启动: ${result.message}`);
      setIsTraining(true);
      
    } catch (error) {
      addLog(`❌ 训练启动失败: ${error.message}`);
    }
  };

  const stopTraining = async () => {
    try {
      addLog('🛑 正在停止训练...');
      
      const response = await fetch('http://localhost:8000/api/training/stop', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        }
      });
      
      const result = await response.json();
      addLog(`✅ 训练停止: ${result.message}`);
      setIsTraining(false);
      
    } catch (error) {
      addLog(`❌ 训练停止失败: ${error.message}`);
    }
  };

  const exportData = async () => {
    try {
      const response = await fetch('http://localhost:8000/api/training/export');
      const result = await response.json();
      addLog(`📤 ${result.message}`);
    } catch (error) {
      addLog(`❌ 数据导出失败: ${error.message}`);
    }
  };

  const clearLogs = () => {
    setLogs([]);
  };

  // 获取性能颜色类名
  const getPerformanceClass = (performance) => {
    if (performance > 0.7) return 'performance-good';
    if (performance > 0.4) return 'performance-medium';
    return 'performance-poor';
  };

  return (
    <div className="App">
      <header className="App-header">
        <h1>翁法罗斯传奇十二因子</h1>
        <p>AI驱动的多智能体训练监控系统</p>
      </header>
      
      {/* 控制面板 */}
      <div className="controls-panel">
        <button 
          className={`btn ${isTraining ? 'btn-danger' : 'btn-primary'}`}
          onClick={isTraining ? stopTraining : startTraining}
        >
          {isTraining ? `停止训练 (${trainingStatus.progress.toFixed(1)}%)` : '开始训练'}
        </button>
        <button className="btn btn-secondary" onClick={exportData}>
          导出数据
        </button>
        <button className="btn btn-secondary" onClick={clearLogs}>
          清空日志
        </button>
      </div>

      <div className="dashboard">
        {/* 系统状态面板 */}
        <div className="status-panel">
          <h2>系统状态</h2>
          <div className="status-grid">
            <div className="status-item">
              <div>当前轮回</div>
              <div className="status-value">{gameStatus.cycle}</div>
            </div>
            <div className="status-item">
              <div>铁幕进度</div>
              <div className="status-value">{gameStatus.iron_curtain.toFixed(1)}%</div>
            </div>
            <div className="status-item">
              <div>训练状态</div>
              <div className="status-value">{gameStatus.training_status}</div>
            </div>
            <div className="status-item">
              <div>活跃连接</div>
              <div className="status-value">{gameStatus.active_connections}</div>
            </div>
          </div>
        </div>

        {/* 训练性能面板 */}
        <div className="status-panel">
          <h2>训练性能</h2>
          <div className="status-grid">
            <div className="status-item">
              <div>总体性能</div>
              <div className={`status-value ${getPerformanceClass(trainingMetrics.overall_performance)}`}>
                {(trainingMetrics.overall_performance * 100).toFixed(1)}%
              </div>
            </div>
            <div className="status-item">
              <div>平均奖励</div>
              <div className="status-value">{trainingMetrics.average_reward.toFixed(2)}</div>
            </div>
            <div className="status-item">
              <div>生存率</div>
              <div className="status-value">{(trainingMetrics.survival_rate * 100).toFixed(1)}%</div>
            </div>
            <div className="status-item">
              <div>训练损失</div>
              <div className="status-value">{trainingMetrics.training_loss.toFixed(4)}</div>
            </div>
          </div>
        </div>

        {/* 训练进度面板 */}
        <div className="status-panel">
          <h2>训练进度</h2>
          <div className="progress-info">
            <div>周期: {trainingStatus.current_cycle} / {trainingStatus.total_cycles}</div>
            <div className="progress-bar">
              <div 
                className="progress-fill"
                style={{width: `${trainingStatus.progress}%`}}
              ></div>
            </div>
            <div>进度: {trainingStatus.progress.toFixed(1)}%</div>
          </div>
        </div>

        {/* 十二因子状态 */}
        <div className="factors-panel">
          <h2>十二因子状态</h2>
          <div className="factors-grid">
            {activeFactors.map((factor, index) => (
              <div key={index} className="factor-card">
                <div className="factor-header">
                  <div className="factor-name">{factor.name}</div>
                  <div className="factor-type">{factor.type}</div>
                </div>
                <div>等级: {factor.level}</div>
                <div>位置: ({factor.x.toFixed(1)}, {factor.y.toFixed(1)})</div>
                <div className="health-bar">
                  <div 
                    className="health-fill"
                    style={{width: `${(factor.health / 200) * 100}%`}}
                  ></div>
                </div>
                <div>生命: {factor.health.toFixed(1)}/200</div>
                <div>状态: {factor.alive ? '✅ 存活' : '❌ 阵亡'}</div>
              </div>
            ))}
          </div>
        </div>

        {/* 实时日志 */}
        <div className="logs-panel">
          <h2>实时训练日志</h2>
          <div className="log-output">
            {logs.length > 0 ? (
              logs.map((log, index) => (
                <div key={index} className="log-entry">
                  <span className="log-time">{log.time}</span>: {log.message}
                </div>
              ))
            ) : (
              <div className="log-entry">等待训练开始...</div>
            )}
          </div>
        </div>
      </div>
    </div>
  );
}

export default App;